package com.ydl.learning.flink.demo.broadcast

import com.alibaba.fastjson.JSON
import com.alibaba.fastjson.serializer.SerializeFilter
import com.ydl.learning.flink.demo.Utils
import org.apache.flink.api.common.state.MapStateDescriptor
import org.apache.flink.api.common.typeinfo.{BasicTypeInfo, TypeHint, TypeInformation}
import org.apache.flink.api.scala._
import org.apache.flink.streaming.api.functions.co.KeyedBroadcastProcessFunction
import org.apache.flink.util.Collector
import org.slf4j.{Logger, LoggerFactory}
import scala.collection.JavaConversions._

/**
 * 广播流demo
 * 广播流分keyed流和non-keyed流，本例是non-keyed流
 *
 * @author ydl
 * @since 2021/5/13
 */
object BroadcastStreamDemo2 extends App with Utils {
  private val log: Logger = LoggerFactory.getLogger("BroadcastStreamDemo2")
  val info = getSocketTextStream(9999)
  val ruleStream = info._1.socketTextStream("localhost", 8888, '\n')
      .map(data => new Rule(data, data))
  //广播描述
  val ruleStateDescriptor: MapStateDescriptor[String, Rule] = new MapStateDescriptor[String, Rule](
    "RulesBroadcastState",
    BasicTypeInfo.STRING_TYPE_INFO,
    TypeInformation.of(new TypeHint[Rule]() {})
  )
  val ruleBroadcastStream = ruleStream.broadcast(ruleStateDescriptor)

  val colorPartitionedStream = info._2
      .map(data => {
        val color = new Color(data)
        val shape = new Shape(data)
        new Item(color, shape)
      })
      .keyBy(_.getValue)

  val output = colorPartitionedStream.connect(ruleBroadcastStream).process(
    // KeyedBroadcastProcessFunction 中的类型参数表示：
    //   1. key stream 中的 key 类型
    //   2. 非广播流中的元素类型
    //   3. 广播流中的元素类型
    //   4. 结果的类型，在这里是 string
    new KeyedBroadcastProcessFunction[Color, Item, Rule, String] {

      // 存储部分匹配的结果，即匹配了一个元素，正在等待第二个元素
      // 我们用一个数组来存储，因为同时可能有很多第一个元素正在等待
      private val mapStateDesc: MapStateDescriptor[String, List[Item]] = new MapStateDescriptor[String, List[Item]](
        "items",
        BasicTypeInfo.STRING_TYPE_INFO,
        createTypeInformation[List[Item]])
      //new ListTypeInfo[Item](classOf[Item]))
      // 与之前的 ruleStateDescriptor 相同
      private val ruleStateDescriptor_ : MapStateDescriptor[String, Rule] = ruleStateDescriptor

      override def processElement(value: Item,
                                  ctx: KeyedBroadcastProcessFunction[Color, Item, Rule, String]#ReadOnlyContext,
                                  out: Collector[String]): Unit = {
        log.info("get data:{}", JSON.toJSONString(value, new Array[SerializeFilter](0)))
        val state = getRuntimeContext.getMapState(mapStateDesc)
        val shape = value.getShape

        for (entry <- ctx.getBroadcastState(ruleStateDescriptor_).immutableEntries) {
          val ruleName = entry.getKey
          val rule = entry.getValue
          var stored = state.get(ruleName)
          if (stored == null) stored = List[Item]()
          if (shape.getValue.equals(rule.getSecond) && !stored.isEmpty) {
            for (i <- stored) {
              out.collect("MATCH: " + i + " - " + value)
            }
            stored.clear()
          }
          // 不需要额外的 else{} 段来考虑 rule.first == rule.second 的情况
          if (shape.getValue.equals(rule.getFirst)) stored = stored :+ value
          if (stored.isEmpty) state.remove(ruleName)
          else state.put(ruleName, stored)
        }
      }

      override def processBroadcastElement(value: Rule,
                                           ctx: KeyedBroadcastProcessFunction[Color, Item, Rule, String]#Context,
                                           out: Collector[String]): Unit = {
        log.info("get broadcast:{}", JSON.toJSONString(value, new Array[SerializeFilter](0)))
        ctx.getBroadcastState(ruleStateDescriptor_).put(value.getFirst, value)
      }
    })
  //colorPartitionedStream.addSink(new LogSink[String]).name("log")
  info._1.execute("BroadcastStreamDemo")
}
